from tensorflow import keras
from tensorflow.keras import datasets  # 导入经典数据集加载模块
import tensorflow as tf  # 导入 TF 库


def preprocess(x, y):
    x = tf.cast(x, dtype=tf.float32) / 255.
    x = tf.reshape(x, [-1, 28 * 28])
    y = tf.cast(y, dtype=tf.int32)
    y = tf.one_hot(y, depth=10)
    return x, y


if __name__ == '__main__':
    (x, y), (x_test, y_test) = datasets.mnist.load_data()
    print('x:', x.shape, 'y:', y.shape, 'x test:', x_test.shape, 'y test:',
          y_test)

    train_db = tf.data.Dataset.from_tensor_slices((x, y))  # 构建 Dataset 对象

    train_db = train_db.shuffle(10000)# 随机打散样本，不会打乱样本与标签映射关系
    train_db = train_db.batch(128)  # 设置批训练，batch size 为 128
    train_db = train_db.map(preprocess)# 预处理函数实现在 preprocess 函数中，传入函数名即可


hh


